# coding=utf-8
import math
import numpy as np
import tensorflow as tf
slim = tf.contrib.slim
import a11_net_params as netparam
import a13_gen_anchor_box
"""
对样本及对应的物体的bbox进行预处理，这个过程称为Encode, 输入为样本的类别和bbox
输出为特征金字塔各层对应的每个prior box的真实类别以及偏移
Encode labels and bounding boxes.
labels     - 样本对应的类别标签
bboxes     - 样本中各个目前物体的bounding box
anchors    - 网络中各层feature maps对应的prior boxes
"""
def bboxes_encode(labels, bboxes, anchors, scope=None):
    return tf_ssd_bboxes_encode(
    	labels, 
		bboxes, 
		anchors,
		netparam.default_params.num_classes,
		netparam.default_params.no_annotation_label,
		ignore_threshold = 0.5,
		prior_scaling = netparam.default_params.prior_scaling,
		scope=scope)

"""
Encode groundtruth labels and bounding boxes using SSD net anchors.
Encoding boxes for all feature layers.
Arguments:
  labels: 1D Tensor(int64) containing groundtruth labels;
  bboxes: Nx4 Tensor(float) with bboxes relative coordinates;
  anchors: List of Numpy array with layer anchors;
  matching_threshold: Threshold for positive match with groundtruth bboxes;
  prior_scaling: Scaling of encoded coordinates.
Return:
  (target_labels, target_localizations, target_scores):
    Each element is a list of target Tensors.
"""
def tf_ssd_bboxes_encode(labels,
                         bboxes,
                         anchors,
                         num_classes,
                         no_annotation_label,
                         ignore_threshold=0.5,
                         prior_scaling=[0.1, 0.1, 0.2, 0.2],
                         dtype=tf.float32,
                         scope='ssd_bboxes_encode'):

    with tf.name_scope(scope):
        target_labels = []
        target_localizations = []
        target_scores = []
        for i, anchors_layer in enumerate(anchors): # each pyramid feature maps
            with tf.name_scope('bboxes_encode_block_%i' % i):
                t_labels, t_loc, t_scores = tf_ssd_bboxes_encode_layer(
                    labels, 
                    bboxes, 
                    anchors_layer,
                    num_classes, 
                    no_annotation_label,
                    ignore_threshold,
                    prior_scaling, 
                    dtype)
                target_labels.append(t_labels)
                target_localizations.append(t_loc)
                target_scores.append(t_scores)
        return target_labels, target_localizations, target_scores

"""
针对一层的feature maps对ground truth进行操作
Encode groundtruth labels and bounding boxes using SSD anchors from
one layer.

Arguments:
  labels: 1D Tensor(int64) containing groundtruth labels;
  bboxes: Nx4 Tensor(float) with bboxes relative coordinates;
  anchors_layer: Numpy array with layer anchors;
  matching_threshold: Threshold for positive match with groundtruth bboxes;
  prior_scaling: Scaling of encoded coordinates.

Return:
  (target_labels, target_localizations, target_scores): Target Tensors.
"""
def tf_ssd_bboxes_encode_layer(labels,
                               bboxes,
                               anchors_layer,
                               num_classes,
                               no_annotation_label,
                               ignore_threshold=0.5,
                               prior_scaling=[0.1, 0.1, 0.2, 0.2],
                               dtype=tf.float32):
    # Anchors coordinates and volume.
    # 将中心点 + size的表示方法展开成左上角坐标+右下角坐标
    # 以第一个feature map为例：xref和yref的shape为(38,38,1), href和wref的shape为
    # (4),则展开后ymin,xmin,ymax,xmax的shape为(38,38,4)
    yref, xref, href, wref = anchors_layer
    ymin = yref - href / 2.
    xmin = xref - wref / 2.
    ymax = yref + href / 2.
    xmax = xref + wref / 2.
    vol_anchors = (xmax - xmin) * (ymax - ymin)

    # Initialize tensors...
    # 定义一些变量，后边用到
    shape = (yref.shape[0], yref.shape[1], href.size)
    feat_labels = tf.zeros(shape, dtype=tf.int64)
    feat_scores = tf.zeros(shape, dtype=dtype)
    feat_ymin = tf.zeros(shape, dtype=dtype)
    feat_xmin = tf.zeros(shape, dtype=dtype)
    feat_ymax = tf.ones(shape, dtype=dtype)
    feat_xmax = tf.ones(shape, dtype=dtype)
    
    def jaccard_with_anchors(bbox):
        """Compute jaccard score between a box and the anchors.
        """
        int_ymin = tf.maximum(ymin, bbox[0])
        int_xmin = tf.maximum(xmin, bbox[1])
        int_ymax = tf.minimum(ymax, bbox[2])
        int_xmax = tf.minimum(xmax, bbox[3])
        h = tf.maximum(int_ymax - int_ymin, 0.)
        w = tf.maximum(int_xmax - int_xmin, 0.)
        # Volumes.
        # 交集面积
        inter_vol = h * w 
        # 并集面积
        union_vol = vol_anchors - inter_vol + (bbox[2] - bbox[0]) * (bbox[3] - bbox[1])
        # Jaccard值 = 两者交集面积 / 两者并集面积
        jaccard = tf.div(inter_vol, union_vol)
        return jaccard

    def intersection_with_anchors(bbox):
        """Compute intersection between score a box and the anchors.
        """
        int_ymin = tf.maximum(ymin, bbox[0])
        int_xmin = tf.maximum(xmin, bbox[1])
        int_ymax = tf.minimum(ymax, bbox[2])
        int_xmax = tf.minimum(xmax, bbox[3])
        h = tf.maximum(int_ymax - int_ymin, 0.)
        w = tf.maximum(int_xmax - int_xmin, 0.)
        inter_vol = h * w
        scores = tf.div(inter_vol, vol_anchors)
        return scores

    def condition(i, feat_labels, feat_scores, feat_ymin, feat_xmin, feat_ymax, feat_xmax):
        """Condition: check label index.
        """
        r = tf.less(i, tf.shape(labels))
        return r[0]

    def body(i, feat_labels, feat_scores, feat_ymin, feat_xmin, feat_ymax, feat_xmax):
        """Body: update feature labels, scores and bboxes.
        Follow the original SSD paper for that purpose:
          - assign values when jaccard > 0.5;
          - only update if beat the score of other bboxes.
        """
        # Jaccard score.
        label = labels[i]
        bbox = bboxes[i]
        # 返回当前ground truth与每一个anchor box的交并比
        jaccard = jaccard_with_anchors(bbox)
        
        # Mask: check threshold + scores + no annotations + num_classes.
        # 筛选掉交并比 < 0的
        mask = tf.greater(jaccard, feat_scores)
        # mask = tf.logical_and(mask, tf.greater(jaccard, matching_threshold))
        mask = tf.logical_and(mask, feat_scores > -0.5)
        mask = tf.logical_and(mask, label < num_classes)
        imask = tf.cast(mask, tf.int64)
        fmask = tf.cast(mask, dtype)
        # Update values using mask.
        feat_labels = imask * label + (1 - imask) * feat_labels
        feat_scores = tf.where(mask, jaccard, feat_scores)

        feat_ymin = fmask * bbox[0] + (1 - fmask) * feat_ymin
        feat_xmin = fmask * bbox[1] + (1 - fmask) * feat_xmin
        feat_ymax = fmask * bbox[2] + (1 - fmask) * feat_ymax
        feat_xmax = fmask * bbox[3] + (1 - fmask) * feat_xmax

        # Check no annotation label: ignore these anchors...
        # interscts = intersection_with_anchors(bbox)
        # mask = tf.logical_and(interscts > ignore_threshold,
        #                       label == no_annotation_label)
        # # Replace scores by -1.
        # feat_scores = tf.where(mask, -tf.cast(mask, dtype), feat_scores)

        return [i+1, feat_labels, feat_scores, feat_ymin, feat_xmin, feat_ymax, feat_xmax]
    
    # Main loop definition.
    # tf.while_loop()如果满足condition，就执行body，后边list为参数
    i = 0
    [i, feat_labels, feat_scores, feat_ymin, feat_xmin, feat_ymax, feat_xmax] = \
    tf.while_loop(condition, body,
    [i, feat_labels, feat_scores, feat_ymin, feat_xmin, feat_ymax, feat_xmax])
    
    # Transform to center / size.
    feat_cy = (feat_ymax + feat_ymin) / 2.
    feat_cx = (feat_xmax + feat_xmin) / 2.
    feat_h = feat_ymax - feat_ymin
    feat_w = feat_xmax - feat_xmin
    # Encode features.
    feat_cy = (feat_cy - yref) / href / prior_scaling[0]
    feat_cx = (feat_cx - xref) / wref / prior_scaling[1]
    feat_h = tf.log(feat_h / href) / prior_scaling[2]
    feat_w = tf.log(feat_w / wref) / prior_scaling[3]
    # Use SSD ordering: x / y / w / h instead of ours.
    feat_localizations = tf.stack([feat_cx, feat_cy, feat_w, feat_h], axis=-1)
    return feat_labels, feat_localizations, feat_scores

"""
Demo
"""
def demo_tf_ssd_bboxes_encode_layer():
	labels = np.array([[0,1,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0]])
	bboxes = np.array([[5,5, 50, 80]])
	anchors_layer = a13_gen_anchor_box.anchors()
	tf_ssd_bboxes_encode_layer(
		labels,
		bboxes,
		anchors_layer[0],
		netparam.default_params.num_classes,
		netparam.default_params.no_annotation_label,
		0.5,
		netparam.default_params.prior_scaling,
		tf.float32)

if __name__ == '__main__':
	demo_tf_ssd_bboxes_encode_layer()
	print 'OK'

